package com.senke.aiagent.agent;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具调用的基础代理类，具体实现了 think 和 act 方法，可以用作创建实例的父类
 */
@Slf4j
@Data
@EqualsAndHashCode(callSuper = true)
public class ToolCallAgent extends ReActAgent {

    // 可用的工具列表
    private final ToolCallback[] availableTools;

    // 工具调用管理器
    private final ToolCallingManager toolCallingManager;

    // 工具调用响应
    private ChatResponse toolCallChatResponse;

    // 聊天选项，用于控制聊天行为
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        // 禁用內置的工具调用机制，自己维护上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
    }

    /**
     * 处理当前状态并决定下一步行动
     *
     * @return 是否需要执行行动
     */
    @Override
    public boolean think() {
        // 校验提示词是否为空
        if (!StringUtils.isBlank(this.getNextStepPrompt())) {
            getMessageList().add(new UserMessage(this.getNextStepPrompt()));
        }
        Prompt prompt = new Prompt(this.getMessageList(), chatOptions);
        try {
            // 构建工具调用请求
            ChatResponse chatResponse = this.getChatClient().prompt(prompt)
                    .system(this.getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            // 保存工具调用响应
            toolCallChatResponse = chatResponse;
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            // 打印日志
            String result = assistantMessage.getText();
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
            //log.info(this.getName() + "的思考结果：" + result);
            log.info(this.getName() + "使用了{}个工具", toolCalls.size());
            String toolCallInfo = toolCalls.stream()
                    .map(toolCall -> String.format("工具名称：%s，参数：%s",
                            toolCall.name(),
                            toolCall.arguments()))
                    .collect(Collectors.joining("\n"))
                    .toString();
            log.info(this.getName() + "工具调用信息：\n{}", toolCallInfo);
            // 不需要调用工具时，记录助手消息
            if (toolCalls.isEmpty()) {
                this.getMessageList().add(assistantMessage);
                return false;
            }
            return true;
        } catch (Exception e) {
            log.error(this.getName() + "思考出现问题：" + e.getMessage());
            this.getMessageList().add(new AssistantMessage("思考出现问题：" + e.getMessage()));
            return false;
        }
    }

    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "无工具调用";
        }
        // 调用工具
        Prompt prompt = new Prompt(this.getMessageList(), chatOptions);
        List<Message> messages = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse).conversationHistory();
        // 记录消息上下文，包括助手消息和工具调用结果
        this.setMessageList(messages);
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) messages.getLast();
        String result = toolResponseMessage.getResponses().stream()
                .map(response -> String.format("工具%s成功调用，结果：%s", response.name(), response.responseData().toString()))
                //.map(response -> String.format("工具%s成功调用，结果：%s", response.name(), " "))
                .collect(Collectors.joining("\n"))
                .toString();
        // 若调用了终止工具，则将状态设置为已完成
        boolean terminateToolCalled = toolResponseMessage.getResponses().stream()
                .anyMatch(response -> "doTerminate".equals(response.name()));
        if (terminateToolCalled) {
            this.setState(AgentState.FINISHED);
            log.info(this.getName() + "已完成任务");
        }
        // 打印日志
        log.info(result);
        return result;
    }
}
